/*#########################################
# Copyright (c) [2010-2011], Josef Robert Novak
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
#  modification, are permitted #provided that the following conditions
#  are met:
#
# * Redistributions of source code must retain the above copyright 
#    notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above 
#    copyright notice, this list of #conditions and the following 
#    disclaimer in the documentation and/or other materials provided 
#    with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS 
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE 
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, 
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) 
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, 
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED 
# OF THE POSSIBILITY OF SUCH DAMAGE.
#########################################*/
#include "htk2wfst.h"
#include <iomanip>


HTK2WFST::HTK2WFST( const char* prefix_ ){
    //Setup output files
    prefix=prefix_; char fname[1024];
    
    sprintf(fname, "%s.st", prefix);
    states_ofp.open(fname);
    sprintf(fname, "%s.hmm", prefix);
    hmms_ofp.open(fname);
    sprintf(fname, "%s.var", prefix);
    vars_ofp.open(fname);
    sprintf(fname, "%s.mean", prefix);
    means_ofp.open(fname);
    sprintf(fname, "%s.pdf", prefix);
    pdfs_ofp.open(fname);
    sprintf(fname, "%s.dur", prefix);
    durs_ofp.open(fname);
    sprintf(fname, "%s.carcs", prefix);
    carcs_ofp.open(fname);
    sprintf(fname, "%s.harcs", prefix);
    harcs_ofp.open(fname);
}

HTK2WFST::~HTK2WFST(){
    //Close output file descriptors
    states_ofp.close();
    means_ofp.close();
    vars_ofp.close();
    durs_ofp.close();
    hmms_ofp.close();
    pdfs_ofp.close();
    carcs_ofp.close();
    harcs_ofp.close();
    //Write symbol tables
    char fname[1024];
    sprintf(fname, "%s.hmm.syms", prefix);
    hmm_syms.WriteTable(fname);
    sprintf(fname, "%s.st.syms", prefix);
    state_syms.WriteTable(fname);
    sprintf(fname, "%s.mean.syms", prefix);
    mean_syms.WriteTable(fname);
    sprintf(fname, "%s.var.syms", prefix);
    var_syms.WriteTable(fname);
}

void HTK2WFST::parse_hmmdefs( const char *hmmdefs, const char* prefix ){
    FILE *hmmdefs_ifp;
    if ( (hmmdefs_ifp = fopen( hmmdefs , "rb" )) == NULL )
      printf("HTKModels::Models(2) - printf opening htkModelsFName") ;

	if ( htkparse( (void *)hmmdefs_ifp ) != 0 )
	  cout << "Parsing failed." << endl;

   output_wfst_am( ) ;

   fclose( hmmdefs_ifp ) ;
}

vector<string> HTK2WFST::split(const string& s, const string& delim, const bool keep_empty) {
    vector<string> result;
    if (delim.empty()) {
        result.push_back(s);
        return result;
    }
    string::const_iterator substart = s.begin(), subend;
    while (true) {
        subend = search(substart, s.end(), delim.begin(), delim.end());
        string temp(substart, subend);
        if (keep_empty || !temp.empty()) {
            result.push_back(temp);
        }
        if (subend == s.end()) {
            break;
        }
        substart = subend + delim.size();
    }
    return result;
}

void HTK2WFST::output_wfst_am( ) {
    tot_states = 1;
    carcs_ofp << htk_def.n_hmms+1 << " 2" << endl;
    harcs_ofp << htk_def.n_sh_states+1 << " 2" << endl;
    for ( int i=0 ; i<htk_def.n_sh_states ; i++ ){
        parse_gmm( htk_def.sh_states[i] ) ;
        output_durations( htk_def.sh_states[i] );
    }

    for ( int i=0 ; i<htk_def.n_hmms ; i++ ){
        hmms_ofp  << hmm_syms.AddSymbol(htk_def.hmms[i]->name) << " 1";
        carcs_ofp << hmm_syms.AddSymbol(htk_def.hmms[i]->name) << " 0 " << htk_def.hmms[i]->n_states-2 << " " << htk_def.hmms[i]->name << endl;
        for ( int j=1 ; j<(htk_def.hmms[i]->n_states-1) ; j++ ){
            if ( htk_def.hmms[i]->emit_states[j-1]->sh_name != NULL ){
                hmms_ofp << " " << state_syms.AddSymbol(htk_def.hmms[i]->emit_states[j-1]->sh_name);
                output_carc( htk_def.hmms[i]->emit_states[j-1] );
            }else{
                cerr << "State not shared, may cause a problem!" << endl;
                parse_gmm( htk_def.hmms[i]->emit_states[j-1] ) ;
            }
        }
        hmms_ofp << endl;
    }
    carcs_ofp << "24702 0 1 EOW" << endl;
    carcs_ofp << "0 0 0" << endl;

    return;
}

void HTK2WFST::output_carc( HTKHMMState *st ) {
    std::string sh_name = st->sh_name;
    vector<string> items = split( sh_name, "_" );
    string tmat_s = items[0].substr(1,1) + "_" + items[1];
    int    i = find_tmat( tmat_s.c_str() );
    int    j = atoi(items[2].c_str())-1;
    if( i >= 0 ){
        double dur1 = htk_def.sh_transmats[i]->transp[j][j];
        double dur2 = 1.0-dur1;
        carcs_ofp << state_syms.AddSymbol(sh_name) << " " << setprecision(15) << dur1 << " " << setprecision(15) << dur2 << endl;
    }else{
        for( int k=0; k<htk_def.n_hmms; k++ ){
            if( strcmp( items[1].c_str(), htk_def.hmms[k]->name ) == 0 ){
                double dur1 = htk_def.hmms[k]->transmat->transp[j][j];
                double dur2 = 1.0-dur1;
                carcs_ofp << state_syms.AddSymbol(sh_name) << " " << setprecision(15) << dur1 << " " << setprecision(15) << dur2 << endl;
            }
        }
    }    
}

void HTK2WFST::output_durations( HTKHMMState *st ) {
    std::string sh_name = st->sh_name;
    vector<string> items = split( sh_name, "_" );
    string tmat_s = items[0].substr(1,1) + "_" + items[1];
    int    i = find_tmat( tmat_s.c_str() );
    int    j = atoi(items[2].c_str())-1;
    
    harcs_ofp << state_syms.AddSymbol(sh_name) << " 0 x" << endl;
    if( i >= 0 ){
        double tr1  = htk_def.sh_transmats[i]->transp[j][j];
        double tr2  = 1.0-tr1;
        double dur1 = (1.0/(1.0-tr1));
        double dur2 = dur1 * dur1 * 0.5;
        
        durs_ofp  << state_syms.AddSymbol(sh_name) << " 1 " << setprecision(15) << dur1 << " " << setprecision(15) << dur2 << endl;
        harcs_ofp << state_syms.AddSymbol(sh_name) << " " << setprecision(15) << tr1  << " " << setprecision(15) << tr2  << endl;
    }else{
        for( int k=0; k<htk_def.n_hmms; k++ ){
            if( strcmp( items[1].c_str(), htk_def.hmms[k]->name ) == 0 ){
                double tr1  = htk_def.hmms[k]->transmat->transp[j][j];
                double tr2  = 1.0-tr1;
                double dur1 = (1.0/(1.0-tr1));
                double dur2 = dur1 * dur1 * 0.5;
                
                durs_ofp  << state_syms.AddSymbol(sh_name) << " 1 " << setprecision(15) << dur1 << " " << setprecision(15) << dur2 << endl;
                harcs_ofp << state_syms.AddSymbol(sh_name) << " " << setprecision(15) << tr1  << " " << setprecision(15) << tr2  << endl;
            }
        }
    }    
}

void HTK2WFST::parse_gmm( HTKHMMState *st ){
    for( int i=0; i<st->n_mixes; i++ ){
        output_means( st->sh_name, i, st->mixes[i]->means );
        output_vars(  st->sh_name, i, st->mixes[i]->vars );
    }
    
    if( st->mixes != NULL ){
        states_ofp << state_syms.AddSymbol(st->sh_name) << " 1";

        for ( int i=0 ; i<st->n_mixes ; i++ ){
            htk_states.insert(StateMap::value_type(tot_states,st->mixes[i]->weight));
            states_ofp << " " << tot_states << " " << setprecision(15) << st->mixes[i]->weight;
            pdfs_ofp << tot_states << " " << tot_states << " " << tot_states << " 0" << endl;
            tot_states++;
        }
        states_ofp << endl;
    }

    return;
}

void HTK2WFST::output_means( const char *sh_name , int i, float *means ) {
    char mname[1024];
    sprintf(mname, "%s_PDF_%d_MEAN_%d", sh_name, i+1, i+1);
    means_ofp << mean_syms.AddSymbol(mname) << " 1 ";
    for( int j=0; j<htk_def.global_opts.vec_size; j++)
        means_ofp << setprecision(15) << means[j] << " ";
    means_ofp << endl;
    return;
}

void HTK2WFST::output_vars( const char *sh_name , int i, float *vars ){
    char vname[1024];
    sprintf(vname, "%s_PDF_%d_VAR_%d", sh_name, i+1, i+1);
    vars_ofp << var_syms.AddSymbol(vname) << " 1 ";
    for ( int j=0 ; j<htk_def.global_opts.vec_size; j++ )
        vars_ofp << setprecision(15) << vars[j] << " ";
    vars_ofp << endl;
    return;
}

int HTK2WFST::find_tmat( const char *sh_name ){
   for ( int i=0 ; i<htk_def.n_sh_transmats ; i++ ){
      if ( strcmp( sh_name, htk_def.sh_transmats[i]->sh_name ) == 0 )
         return i ;
   }
   return -1 ;
}

